"""会议安排任务的数据处理脚本"""
import json
import random
import pandas as pd
import argparse
from pathlib import Path


def generate_availability(persons, seed=None):
    """
    生成人员的日程可用性数据
    
    Args:
        persons: 人员列表
        seed: 随机种子
        
    Returns:
        可用性字典
    """
    if seed is not None:
        random.seed(seed)
    
    days = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']
    time_slots = ['morning', 'afternoon', 'evening']
    
    availability = {}
    
    for person in persons:
        availability[person] = {}
        for day in days:
            availability[person][day] = {}
            for slot in time_slots:
                # 70%概率空闲，30%概率已安排
                status = 'available' if random.random() > 0.3 else 'busy'
                availability[person][day][slot] = status
    
    return availability


def generate_historical_pattern(meeting_topic, seed=None):
    """
    生成历史会议模式数据
    
    Args:
        meeting_topic: 会议主题
        seed: 随机种子
        
    Returns:
        历史模式字典
    """
    if seed is not None:
        random.seed(seed)
    
    days = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday']
    time_slots = ['morning', 'afternoon', 'evening']
    
    # 随机选择一个常见时间
    most_common_day = random.choice(days)
    most_common_slot = random.choice(time_slots)
    frequency = random.randint(60, 90)  # 60-90%的频率
    
    # 生成最近5次会议记录
    recent_meetings = []
    for week in range(1, 6):
        # 80%概率使用常见时间，20%概率随机
        if random.random() < frequency / 100.0:
            day = most_common_day
            slot = most_common_slot
        else:
            day = random.choice(days)
            slot = random.choice(time_slots)
        
        recent_meetings.append({
            'week': -week,
            'day': day,
            'slot': slot
        })
    
    return {
        'most_common_day': most_common_day,
        'most_common_slot': most_common_slot,
        'frequency': frequency,
        'recent_meetings': recent_meetings
    }


def generate_meeting_task(task_id, seed=None):
    """
    生成一个会议安排任务
    
    Args:
        task_id: 任务ID
        seed: 随机种子
        
    Returns:
        任务数据字典
    """
    if seed is not None:
        random.seed(seed + task_id)
    
    # 会议主题候选
    meeting_topics = [
        "产品规划会",
        "技术评审会",
        "季度总结会",
        "客户对接会",
        "团队建设会",
        "需求讨论会",
        "进度同步会",
        "战略研讨会"
    ]
    
    # 人员名字池
    person_names = [
        "李明", "王芳", "陈静", "刘强", "赵敏",
        "张伟", "孙丽", "周涛", "吴娜", "郑勇",
        "钱红", "黄鹏", "杨雪", "徐波", "朱琳"
    ]
    
    # 随机选择会议主题
    meeting_topic = random.choice(meeting_topics)
    
    # 随机选择2-4个必选人员
    num_required = random.randint(2, 4)
    all_persons = random.sample(person_names, min(num_required + random.randint(2, 5), len(person_names)))
    required_attendees = all_persons[:num_required]
    optional_attendees = all_persons[num_required:]
    
    # 生成日程可用性
    availability = generate_availability(all_persons, seed=seed + task_id + 1000)
    
    # 生成历史模式
    historical_pattern = generate_historical_pattern(meeting_topic, seed=seed + task_id + 2000)
    
    return {
        'task_id': f"meeting_{task_id:04d}",
        'meeting_topic': meeting_topic,
        'required_attendees': required_attendees,
        'optional_attendees': optional_attendees,
        'availability': availability,
        'historical_pattern': historical_pattern,
        # 框架必需字段
        'agent_name': 'meeting_agent',  # 指定使用 MeetingAgentLoop
        'reward_model': {
            'ground_truth': ''  # 会议安排任务不需要ground truth，使用过程奖励
        }
    }


def create_dataset(num_samples, output_path, seed=42):
    """
    创建数据集
    
    Args:
        num_samples: 样本数量
        output_path: 输出文件路径
        seed: 随机种子
    """
    print(f"正在生成 {num_samples} 个会议安排任务...")
    
    tasks = []
    for i in range(num_samples):
        task = generate_meeting_task(i, seed=seed)
        tasks.append(task)
    
    # 转换为DataFrame
    df = pd.DataFrame(tasks)
    
    # 保存为parquet格式
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    df.to_parquet(output_path, index=False)
    print(f"数据集已保存到: {output_path}")
    print(f"数据集大小: {len(df)} 条")
    
    # 保存前3个样本为JSON格式（用于调试）
    json_output_path = output_path.with_suffix('.sample.json')
    sample_data = tasks[:min(3, len(tasks))]
    with open(json_output_path, 'w', encoding='utf-8') as f:
        json.dump(sample_data, f, ensure_ascii=False, indent=2)
    print(f"样本JSON已保存到: {json_output_path}")
    
    # 打印第一个样本的关键字段（用于快速检查）
    print("\n第一个样本的关键字段检查:")
    if tasks:
        first_sample = tasks[0]
        print(f"  task_id: {first_sample.get('task_id', 'MISSING')}")
        print(f"  meeting_topic: {first_sample.get('meeting_topic', 'MISSING')}")
        print(f"  agent_name: '{first_sample.get('agent_name', 'MISSING')}'")
        print(f"  reward_model: {first_sample.get('reward_model', 'MISSING')}")
        if 'reward_model' in first_sample and isinstance(first_sample['reward_model'], dict):
            print(f"    └─ ground_truth: '{first_sample['reward_model'].get('ground_truth', 'MISSING')}'")
        print(f"  required_attendees 数量: {len(first_sample.get('required_attendees', []))}")
        print(f"  optional_attendees 数量: {len(first_sample.get('optional_attendees', []))}")
    
    # 打印一些统计信息
    print("\n数据集统计:")
    print(f"- 平均必选人员数: {df['required_attendees'].apply(len).mean():.2f}")
    print(f"- 平均可选人员数: {df['optional_attendees'].apply(len).mean():.2f}")
    print(f"- 会议主题分布:")
    topic_counts = {}
    for topics in df['meeting_topic']:
        topic_counts[topics] = topic_counts.get(topics, 0) + 1
    for topic, count in sorted(topic_counts.items(), key=lambda x: -x[1]):
        print(f"  {topic}: {count}")


def main():
    parser = argparse.ArgumentParser(description='生成会议安排任务数据集')
    parser.add_argument('--train_size', type=int, default=100, help='训练集大小')
    parser.add_argument('--val_size', type=int, default=20, help='验证集大小')
    parser.add_argument('--output_dir', type=str, default='./meeting_dataset', help='输出目录')
    parser.add_argument('--seed', type=int, default=42, help='随机种子')
    
    args = parser.parse_args()
    
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 生成训练集
    print("=" * 60)
    print("生成训练集")
    print("=" * 60)
    train_path = output_dir / 'train.parquet'
    create_dataset(args.train_size, train_path, seed=args.seed)
    
    # 生成验证集
    print("\n" + "=" * 60)
    print("生成验证集")
    print("=" * 60)
    val_path = output_dir / 'test.parquet'
    create_dataset(args.val_size, val_path, seed=args.seed + 10000)
    
    print("\n" + "=" * 60)
    print("数据生成完成！")
    print("=" * 60)
    print(f"训练集: {train_path}")
    print(f"验证集: {val_path}")


if __name__ == "__main__":
    main()

